from rdkit import Chem
from rdkit.Chem import Draw, AllChem, Descriptors, rdMolTransforms

#####################################################################################
############################## ringed-smiles gen ####################################
#####################################################################################

def calculate_star_distance(mol):
    """
    计算分子中两个星号之间的距离（最短路径长度）。
    :param smiles: 分子的 SMILES 字符串
    :return: 两个星号之间的最短路径长度
    """
    # 查找所有星号的位置
    stars = [atom for atom in mol.GetAtoms() if atom.GetSymbol() == "*"]
    if len(stars) != 2:
        raise ValueError("SMILES 中必须恰好有两个星号")

    # 获取两个星号的原子索引
    star1_idx = stars[0].GetIdx()
    star2_idx = stars[1].GetIdx()

    # 计算两个星号之间的最短路径长度
    path_length = Chem.GetShortestPath(mol, star1_idx, star2_idx)
    return len(path_length) - 1  # 减1是因为路径长度包括起始和结束原子

def create_ring_from_psmiles(mol):
    # 查找两个星号的位置
    stars = [atom for atom in mol.GetAtoms() if atom.GetSymbol() == "*"]
    if len(stars) != 2:
        raise ValueError("SMILES 中必须恰好有两个星号")

    # 获取两个星号分别连接的原子
    star1, star2 = stars[0], stars[1]
    star1_neighbors = star1.GetNeighbors()
    star2_neighbors = star2.GetNeighbors()

    if len(star1_neighbors) != 1 or len(star2_neighbors) != 1:
        raise ValueError("每个星号必须恰好连接一个原子")

    # 获取连接星号的原子的索引
    atom1_idx = star1_neighbors[0].GetIdx()
    atom2_idx = star2_neighbors[0].GetIdx()

    # 创建一个可编辑的分子对象
    emol = Chem.EditableMol(mol)

    # 添加单键连接两个原子
    emol.AddBond(atom1_idx, atom2_idx, order=Chem.BondType.SINGLE)

    # 移除星号原子
    # 注意：移除第一个星号后，第二个星号的索引会变化
    emol.RemoveAtom(star1.GetIdx())
    # 重新计算第二个星号的索引（如果第一个星号的索引小于第二个星号的索引）
    if star1.GetIdx() < star2.GetIdx():
        emol.RemoveAtom(star2.GetIdx() - 1)  # 索引减1
    else:
        emol.RemoveAtom(star2.GetIdx())

    # 将可编辑的分子对象转换回普通的 Mol 对象
    new_mol = emol.GetMol()

    # 将修改后的分子转换回 SMILES
    new_smiles = Chem.MolToSmiles(new_mol)
    return new_smiles

def create_large_ring_from_psmiles(psmiles):
    """
    将一个带有星号的 p-SMILES 复制四次，并在星号处连接成键，形成一个大环。
    :param psmiles: 带有星号的 p-SMILES 字符串
    :return: 最终形成的大环的 SMILES 字符串
    """
    # 将 p-SMILES 复制四次，用点分隔
    combined_smiles = ".".join([psmiles] * 4)
    
    # 将连接后的 SMILES 转换为 RDKit 分子对象
    mol = Chem.MolFromSmiles(combined_smiles)
    if mol is None:
        raise ValueError("无效的 SMILES 字符串")
    
    # 查找所有星号的位置
    stars = [atom for atom in mol.GetAtoms() if atom.GetSymbol() == "*"]
    if len(stars) != 8:
        raise ValueError("SMILES 中必须恰好有四个星号")
    
    # 获取每个星号连接的原子
    star_neighbors = [star.GetNeighbors() for star in stars]
    if any(len(neighbors) != 1 for neighbors in star_neighbors):
        raise ValueError("每个星号必须恰好连接一个原子")
    
    # 获取连接星号的原子的索引
    atom_indices = [neighbors[0].GetIdx() for neighbors in star_neighbors]
    
    # 创建一个可编辑的分子对象
    emol = Chem.EditableMol(mol)

    tmp_n = len(atom_indices)
    
    # 添加单键连接星号之间的原子，形成一个环
    for i in range(1,tmp_n,2):
        emol.AddBond(atom_indices[i], atom_indices[(i + 1)%tmp_n], order=Chem.BondType.SINGLE)
    
    # 移除所有星号原子
    for star in sorted(stars, key=lambda x: x.GetIdx(), reverse=True):  # 从索引最大的星号开始移除
        emol.RemoveAtom(star.GetIdx())
    
    # 将可编辑的分子对象转换回普通的 Mol 对象
    new_mol = emol.GetMol()
    
    # 将修改后的分子转换回 SMILES
    new_smiles = Chem.MolToSmiles(new_mol)
    return new_smiles


def ringed_SMILES_gen(input_smiles):
    mol = Chem.MolFromSmiles(input_smiles)
    if mol is None:
        raise ValueError("无效的 SMILES 字符串")
    return (create_ring_from_psmiles(mol), 1) if calculate_star_distance(mol)>6 else (create_large_ring_from_psmiles(input_smiles), 4)


###############################################################################
############################## QSPR models ####################################
###############################################################################

def count_functional_groups(molecule, functional_group):
    if molecule:
        functional_group_count = 0
        functional_group_pattern = Chem.MolFromSmarts(functional_group)
        if functional_group_pattern:
            matches = molecule.GetSubstructMatches(functional_group_pattern)
            functional_group_count = len(matches)
        return functional_group_count
    else:
        return None

#**************labelling smiles********************

def labelling_optical(smiles, state):
    mol = Chem.MolFromSmiles(smiles)
    AllChem.Compute2DCoords(mol)

    # Convert SMILES to RDKit molecule object
    mol = Chem.MolFromSmiles(smiles)

    # Add hydrogen atoms (implicit and explicit)
    mol = Chem.AddHs(mol)

    # Get atom information, excluding hydrogen atoms
    non_hydrogen_atoms = [atom for atom in mol.GetAtoms() if atom.GetSymbol() != 'H']

    # Calculate coordination number for each non-hydrogen atom
    coordination_numbers = [atom.GetDegree() for atom in non_hydrogen_atoms]

    # Print out information about each non-hydrogen atom
    for atom, coordination_number in zip(non_hydrogen_atoms, coordination_numbers):
        element = atom.GetSymbol()

    # Calculate the number of each element with different coordination numbers
    element_coordination_counts = {}

# Selected features
###----atomic coordination number----###
    #---second period---#
    element_coordination_counts['B'] = {}
    element_coordination_counts['B'][3] = 0

    element_coordination_counts['C'] = {}
    element_coordination_counts['C'][4] = 0
    element_coordination_counts['C'][3] = 0
    element_coordination_counts['C'][2] = 0

    element_coordination_counts['N'] = {}
    element_coordination_counts['N'][3] = 0
    element_coordination_counts['N'][2] = 0
    element_coordination_counts['N'][1] = 0

    element_coordination_counts['O'] = {}
    element_coordination_counts['O'][2] = 0
    element_coordination_counts['O'][1] = 0

    element_coordination_counts['F'] = {}
    element_coordination_counts['F'][1] = 0

    #---third period---#
    element_coordination_counts['Si'] = {}
    element_coordination_counts['Si'][4] = 0

    element_coordination_counts['P'] = {}
    element_coordination_counts['P'][6] = 0
    element_coordination_counts['P'][4] = 0
    element_coordination_counts['P'][3] = 0

    element_coordination_counts['S'] = {}
    element_coordination_counts['S'][6] = 0
    element_coordination_counts['S'][4] = 0
    element_coordination_counts['S'][3] = 0
    element_coordination_counts['S'][2] = 0
    element_coordination_counts['S'][1] = 0

    element_coordination_counts['Cl'] = {}
    element_coordination_counts['Cl'][1] = 0

    #---fourth period and beyond---#
    element_coordination_counts['Ge'] = {}
    element_coordination_counts['Ge'][4] = 0

    element_coordination_counts['As'] = {}
    element_coordination_counts['As'][6] = 0
    element_coordination_counts['As'][4] = 0
    element_coordination_counts['As'][3] = 0

    element_coordination_counts['Se'] = {}
    element_coordination_counts['Se'][6] = 0
    element_coordination_counts['Se'][4] = 0
    element_coordination_counts['Se'][3] = 0
    element_coordination_counts['Se'][2] = 0
    element_coordination_counts['Se'][1] = 0

    element_coordination_counts['Br'] = {}
    element_coordination_counts['Br'][1] = 0

    element_coordination_counts['I'] = {}
    element_coordination_counts['I'][1] = 0

    for atom, coordination_number in zip(non_hydrogen_atoms, coordination_numbers):
        element = atom.GetSymbol()
        if element not in element_coordination_counts:
            element_coordination_counts[element] = {}
        if coordination_number not in element_coordination_counts[element]:
            element_coordination_counts[element][coordination_number] = 1
        else:
            element_coordination_counts[element][coordination_number] += 1

###----specific functional groups----###
    # Calculate numbers of rings and aromatic rings
    sssr = Chem.GetSymmSSSR(mol)
    three_ring_count = len([ring for ring in sssr if len(ring) == 3])
    four_ring_count = len([ring for ring in sssr if len(ring) == 4])
    five_ring_count = len([ring for ring in sssr if len(ring) == 5])
    six_ring_count = len([ring for ring in sssr if len(ring) == 6])

    aromatic_count = 0

    for ring in sssr:
        is_aromatic = all(mol.GetAtomWithIdx(atom_index).GetIsAromatic() for atom_index in ring)  
        if is_aromatic:
            aromatic_count += 1

    # Calculate numbers of functional groups
    oh_count = count_functional_groups(mol, '[OH]')
    carbonyl_count = count_functional_groups(mol, 'C=O')
    peroxygen_count = count_functional_groups(mol, 'OO')
    ONC_count = count_functional_groups(mol, 'ON=C')
    CNC_count = count_functional_groups(mol, 'CN=C')
    NNC_count = count_functional_groups(mol, 'NN=C')
    cyanide_count = count_functional_groups(mol, 'C#N')
    isocyanide_count = count_functional_groups(mol, '[C-]#[N+]')
    doublebond_count = count_functional_groups(mol, 'C=C')
    triplebond_count = count_functional_groups(mol, 'C#C')
    CS_count = count_functional_groups(mol, 'C=S')
    SS_count = count_functional_groups(mol, 'S=S')
    benzene_count = count_functional_groups(mol, 'c1ccccc1')
    naphthalene_count = count_functional_groups(mol, 'c1cccc2ccccc12')
    pc1_count = count_functional_groups(mol, 'CC(C)(c1ccccc1)c2ccccc2')
    pc2_count = count_functional_groups(mol, 'c5ccc(C4(c1ccccc1)c2ccccc2c3ccccc34)cc5')
    pma_count = count_functional_groups(mol, 'CCC(OC)=O')

    weight = Descriptors.ExactMolWt(mol)

###----specific functional groups----###
    #features = {}
    features = [
    state,
    element_coordination_counts['B'][3]*10.81/weight,
    element_coordination_counts['C'][4]*12.011/weight,
    element_coordination_counts['C'][3]*12.011/weight,
    element_coordination_counts['C'][2]*12.011/weight,
    element_coordination_counts['N'][3]*14.007/weight,
    element_coordination_counts['N'][2]*14.007/weight,
    element_coordination_counts['N'][1]*14.007/weight,
    element_coordination_counts['O'][2]*16/weight,
    element_coordination_counts['O'][1]*16/weight,
    element_coordination_counts['F'][1]*19/weight,
    element_coordination_counts['Si'][4]*28.085/weight,
    element_coordination_counts['P'][4]*31/weight,
    element_coordination_counts['P'][3]*31/weight,
    element_coordination_counts['S'][4]*32.06/weight,
    element_coordination_counts['S'][3]*32.06/weight,
    element_coordination_counts['S'][2]*32.06/weight,
    element_coordination_counts['S'][1]*32.06/weight,
    element_coordination_counts['Cl'][1]*35.45/weight,
    element_coordination_counts['Ge'][4]*72.63/weight,
    element_coordination_counts['Se'][2]*79/weight,
    element_coordination_counts['Br'][1]*79.904/weight,
    element_coordination_counts['I'][1]*126.9/weight,   # 22 atomic descriptors
    three_ring_count*42/weight,
    four_ring_count*56/weight,
    five_ring_count*70/weight,
    six_ring_count*84/weight,
    # seven_ring_count*98/weight,
    # eight_ring_count*112/weight,
    # large_ring_count*126/weight,
    aromatic_count*78/weight,                        # 8 ring descriptors
    oh_count*17/weight,
    carbonyl_count*28/weight,
    peroxygen_count*32/weight,
    ONC_count*42/weight,
    CNC_count*38/weight,
    NNC_count*40/weight,
    cyanide_count*26/weight,
    isocyanide_count*44/weight,
    doublebond_count*24/weight,
    triplebond_count*24/weight,
    CS_count*44/weight,
    benzene_count*78/weight,
    naphthalene_count*120/weight,
    pc1_count*180/weight, 
    pc2_count*300/weight,
    pma_count*80/weight                            # 16 functional group descriptors
    ]

    return features


    